{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# USAD-LSTM在SWaT数据集上实验"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user/miniconda3/envs/yjq-ml/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np \n",
    "import os\n",
    "import pickle\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from usadlstm_model import  *\n",
    "from utils import  *\n",
    "from evaluator import *\n",
    "import random\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr=1e-4\n",
    "num_epochs=5\n",
    "window_size = 12\n",
    "z_dim = 40\n",
    "batch_size = 5120\n",
    "\n",
    "model_name = \"SWaT\"\n",
    "general_path = \"%s-batch-%s-epochs-%s-hidden-%s\" %(model_name, batch_size, num_epochs, z_dim)\n",
    "\n",
    "model_path = \"../model/USAD-LSTM/\" + general_path + \".pth\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 预处理标准化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_windows(data, window_size):\n",
    "    windows = []\n",
    "    length  = data.shape[0]\n",
    "    for i in range(length):\n",
    "        if length - i >= window_size:\n",
    "            window = data[i:i+window_size]\n",
    "        else:\n",
    "            window = data[i-window_size+1:i+1]\n",
    "        windows.append(window)    \n",
    "    return np.array(windows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import StandardScaler\n",
    "down_sampling_rate = 1\n",
    "train = pd.read_csv(\"../data/SWaT/SWaT_Dataset_Normal_v1.csv\")\n",
    "train = train.drop([\"Timestamp\", \"Normal/Attack\"], axis = 1)\n",
    "for i in list(train):\n",
    "    train[i] = train[i].apply(lambda x: str(x).replace(\",\" , \".\"))\n",
    "train = train.astype(float)\n",
    "train = train.groupby(np.arange(len(train.index)) // down_sampling_rate).mean()\n",
    "min_max_scaler = StandardScaler()\n",
    "train = min_max_scaler.fit_transform(train.values)\n",
    "\n",
    "test = pd.read_csv(\"../data/SWaT/SWaT_Dataset_Attack_v0.csv\",sep=\";\")\n",
    "\n",
    "labels =np.array([ float(label!= 'Normal' ) for label  in test[\"Normal/Attack\"].values])\n",
    "\n",
    "test = test.drop([\"Timestamp\" , \"Normal/Attack\" ] , axis = 1)\n",
    "for i in list(test):\n",
    "    test[i]=test[i].apply(lambda x: str(x).replace(\",\" , \".\"))\n",
    "test = test.astype(float)\n",
    "test = test.groupby(np.arange(len(test.index)) // down_sampling_rate).mean()\n",
    "test = min_max_scaler.transform(test.values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(449919,)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "labels.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(495000, 12, 51)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 转换为时间窗口\n",
    "train_window = convert_to_windows(train, window_size)\n",
    "train_window.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(449919, 12, 51)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_window = convert_to_windows(test, window_size)\n",
    "test_window.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = torch.from_numpy(train_window).float().view([train_window.shape[0], train_window.shape[1]*train_window.shape[2]])\n",
    "test_dataset =  torch.from_numpy(test_window).float().view([test_window.shape[0], test_window.shape[1]*test_window.shape[2]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "train_gaussian_percentage  = 0.2\n",
    "indices = np.random.permutation(len(train_dataset))\n",
    "split_point = int(train_gaussian_percentage * len(train_dataset))\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, drop_last=True,\n",
    "                            sampler=SubsetRandomSampler(indices[:-split_point]), pin_memory=False)\n",
    "val_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, drop_last=True,\n",
    "                        sampler=SubsetRandomSampler(indices[-split_point:]),pin_memory=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [40:46<00:00, 489.36s/it]\n"
     ]
    }
   ],
   "source": [
    "model = UsadLSTM(train[0].shape[0]*window_size, z_dim, lstm_layers = 2)\n",
    "history = model.fit(\n",
    "    train_loader=train_loader,\n",
    "    validation_loader=val_loader,\n",
    "    epochs=num_epochs\n",
    ")\n",
    "\n",
    "torch.save({\n",
    "    'encoder': model.encoder.state_dict(),\n",
    "    'decoder1': model.decoder1.state_dict(),\n",
    "    'decoder2': model.decoder2.state_dict()\n",
    "    }, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test finished\n"
     ]
    }
   ],
   "source": [
    "test_loader = DataLoader(dataset = test_dataset, batch_size=batch_size, shuffle=False, drop_last=False)\n",
    "\n",
    "model = UsadLSTM(test[0].shape[0]*window_size, z_dim, lstm_layers = 2)\n",
    "checkpoint = torch.load(model_path, map_location=get_default_device())\n",
    "model.encoder.load_state_dict(checkpoint['encoder'])\n",
    "model.decoder1.load_state_dict(checkpoint['decoder1'])\n",
    "model.decoder2.load_state_dict(checkpoint['decoder2'])\n",
    "model.to(get_default_device())\n",
    "pred= model.predict_prob(test_loader, alpha = 0.5, beta = 0.5)\n",
    "\n",
    "print(\"test finished\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "results_adjust  = bf_search(labels.squeeze(), pred, verbose=False, is_adjust=True) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'f1-score': 0.8473364521310057,\n",
       " 'precision': 0.8964249232080848,\n",
       " 'recall': 0.803354021200023,\n",
       " 'TP': 43880.0,\n",
       " 'TN': 390228.0,\n",
       " 'FP': 5070.0,\n",
       " 'FN': 10741.0,\n",
       " 'threshold': 9.121429782867398}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results_adjust"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
